package b3_RBTree;

import b2_AVL.AVLTree;

import java.util.ArrayList;
import java.util.List;

/**
 * 红黑树
 */
public class RBTree {

    class RBTreeNode {
        RBTreeNode left = null;
        RBTreeNode right = null;
        RBTreeNode parent = null;
        COLOR color; // 节点的颜色
        int val;

        public RBTreeNode(int val) {
            this.val = val;
            // 默认新增节点为红色
            this.color = COLOR.RED;
        }
    }

    public RBTreeNode root;

    // 插入
    public boolean insert(int val) {
        RBTreeNode node = new RBTreeNode(val);
        if (root == null) {
            this.root = node;
            root.color = COLOR.BLACK;
            return true;
        }

        RBTreeNode parent = null;
        RBTreeNode cur = root;
        while (cur != null) {
            if (val == cur.val) {
                return false;
            } else if (val < cur.val) {
                parent = cur;
                cur = cur.left;
            } else {
                parent = cur;
                cur = cur.right;
            }
        }
        // 此时，cur = null
        if (val < parent.val) {
            parent.left = node;
        } else {
            parent.right = node;
        }
        node.parent = parent;
        cur = node;

        // 调整颜色
        // 新节点插入后，如果parent节点的颜色是红色，一定违反性质三
        while (parent != null && parent.color == COLOR.RED) {
            RBTreeNode grandFather = parent.parent;
            if (parent == grandFather.left) {
                RBTreeNode uncle = grandFather.right;
                if (uncle != null && uncle.color == COLOR.RED) {
                    // 情况一：叔叔节点存在且为红,
                    // 解决方式：将叔叔和父节点改为黑色，祖父节点改为红色
                    // 如果祖父的双亲节点的颜色是红色，需要继续往上调整
                    parent.color = COLOR.BLACK;
                    uncle.color = COLOR.BLACK;
                    grandFather.color = COLOR.RED;
                    // 把 g当成cur，继续向上调整
                    cur = grandFather;
                    parent = cur.parent;
                } else { // 此时，叔叔节点不存在 || 叔叔节点存在，但是颜色是黑色
                    // 再讨论cur是左孩子还是右孩子 ?
                    if (cur == parent.left) {
                        // 情况二
                        rotateR(grandFather);
                        parent.color = COLOR.BLACK;
                        grandFather.color = COLOR.RED;
                    } else {
                        // 情况三
                        rotateL(parent);
                        RBTreeNode temp = parent;
                        parent = cur;
                        cur = temp;
                    }
                }
            } else {
                // parent == grandFather.right
                // 以上情况是插入左边，此时是插入到右边，原理一样
                RBTreeNode uncle = grandFather.left;
                if (uncle != null && uncle.color == COLOR.RED) {
                    // 情况一：叔叔节点存在且为红,
                    // 解决方式：将叔叔和父节点改为黑色，祖父节点改为红色
                    // 如果祖父的双亲节点的颜色是红色，需要继续往上调整
                    parent.color = COLOR.BLACK;
                    uncle.color = COLOR.BLACK;
                    grandFather.color = COLOR.RED;
                    // 把 g当成cur，继续向上调整
                    cur = grandFather;
                    parent = cur.parent;
                } else { // 此时，叔叔节点不存在 || 叔叔节点存在，但是颜色是黑色
                    // 再讨论cur是左孩子还是右孩子 ?
                    if (cur == parent.right) {
                        // 情况二
                        rotateL(grandFather);
                        parent.color = COLOR.BLACK;
                        grandFather.color = COLOR.RED;
                    } else {
                        // 情况三
                        rotateR(parent);
                        RBTreeNode temp = parent;
                        parent = cur;
                        cur = temp;
                    }
                }
            }
        }
        // 在上述循环更新期间，可能会将根节点给成红色，因此此处必须将根节点改为黑色
        root.color = COLOR.BLACK;
        return true;
    }

    // 左单旋
    private void rotateL(RBTreeNode p) {
        // p 的母节点
        RBTreeNode pp = p.parent;
        // p 的右孩子
        RBTreeNode subR = p.right;
        // subR 的左孩子，可能不存在
        RBTreeNode subRL = subR.left;

        // subR 提上去
        if (pp == null) {
            this.root = subR;
        } else if (pp.left == p) {
            pp.left = subR;
        } else {
            // pp.right == parent
            pp.right = subR;
        }
        subR.parent = pp;
        // p 作为 subR 的左孩子
        subR.left = p;
        p.parent = subR;
        // p 与 subRL 连接
        p.right = subRL;
        if (subRL != null) {
            subRL.parent = p;
        }
    }

    // 右单旋
    private void rotateR(RBTreeNode p) {
        // p 的父节点
        RBTreeNode pp = p.parent;
        // p 的左孩子
        RBTreeNode subL = p.left;
        // subL 的右孩子，可能不存在
        RBTreeNode subLR = subL.right;

        if (pp == null) {
            this.root = subL;
        } else if (pp.left == p) {
            pp.left = subL;
        } else {
            pp.right = subL;
        }
        subL.parent = pp;
        subL.right = p;
        p.parent = subL;
        p.left = subLR;
        if (subLR != null) {
            subLR.parent = p;
        }
    }

    /**
     * 打印二叉树, 中序遍历的结果
     **/
    @Override
    public String toString() {
        List<Integer> list = new ArrayList<>();
        inOrder(list, root);
        return list.toString() + "\n" + "是否标准红黑树: " + isValidRBTree();
    }
    // 中序遍历
    private void inOrder(List<Integer> list, RBTreeNode root) {
        if (root == null) {
            return;
        }
        inOrder(list, root.left);
        list.add(root.val);
        inOrder(list, root.right);
    }
    /**
     * 检验是否符合红黑树的性质
     */
    private boolean isValidRBTree() {
        // 空树也是红黑树
        if (root == null) {
            return true;
        }
        // 根节点是黑色
        if (root.color != COLOR.BLACK) {
            System.err.println("违反了性质2：根节点不是黑色");
            return false;
        }
        // 获取单条路径中节点的个数
        int blackCount = 0;
        RBTreeNode cur = root;
        while (cur != null) {
            if (cur.color == COLOR.BLACK) {
                blackCount ++;
            }
            cur = cur.left;
        }
        // 具体的检验方式
        return _isValidRBtree(root, 0, blackCount);
    }
    // 检验是否存在连续的红色节点
    // 检验是否每条路径黑色节点数相同
    private boolean _isValidRBtree(RBTreeNode root, int pathCount, int blackCount) {
        if (root == null) {
            return true;
        }
        // 遇到一个黑色节点，统计当前路径中黑色节点个数
        if(root.color == COLOR.BLACK) {
            pathCount ++;
        }
        // 验证性质4
        RBTreeNode parent = root.parent;
        if(parent != null && parent.color == COLOR.RED && root.color == COLOR.RED) {
            System.err.println("违反了性质4：有连在一起的红色节点");
            return true;
        }
        // 验证性质5
        // 如果是叶子节点，则一条路径已经走到底，检验该条路径中黑色节点总个数是否与先前统计的结果相同
        if (root.left == null && root.right == null) {
            if (pathCount != blackCount) {
                System.err.println("违反了性质5：每条路径中黑色节点个数不一致");
                return false;
            }
        }
        // 以递归的方式检测 root 的左右子树
        return _isValidRBtree(root.left, pathCount, blackCount) &&
                _isValidRBtree(root.right, pathCount, blackCount);
    }
}
